﻿using System;
using System.Collections;
using System.Collections.Generic;
using System.Text.RegularExpressions;
using UnityEngine;
using UnityEngine.Networking;

public class LLMClient : MonoBehaviour
{
    // Singleton
    public static LLMClient Instance;

    // Public events for demo layer
    public event Action<string> OnRequestFailed;  // receives an error message
    public event Action<int, string> OnRateLimit; // (retrySeconds, message)


    [Header("OpenAI Settings (from APIConfig)")]
    private string endpoint => APIConfig.Instance != null ? APIConfig.Instance.chatEndpoint : null;
    private string apiKey => APIConfig.Instance != null ? APIConfig.Instance.chatApiKey : null;
    private string model => APIConfig.Instance != null ? APIConfig.Instance.chatModel : "gpt-4o-mini";

    // Simple DTO that demo code can reuse if desired
    [Serializable]
    public class Message
    {
        public string role;
        public string content;
        public Message(string r, string c)
        {
            role = r;
            content = c;
        }
    }

    // --- Internal structures for JSON serialization / deserialization ---
    [Serializable]
    private class ChatMessage
    {
        public string role;
        public string content;
        public ChatMessage(string r, string c) { role = r; content = c; }
    }

    [Serializable]
    private class ChatRequest
    {
        public string model;
        public ChatMessage[] messages;
    }

    [Serializable]
    private class ResponseMessage
    {
        public string role;
        public string content;
    }

    [Serializable]
    private class Choice
    {
        public ResponseMessage message;
    }

    [Serializable]
    private class RootResponse
    {
        public Choice[] choices;
    }

    // --- Unity lifecycle ---
    void Awake()
    {
        if (Instance == null)
        {
            Instance = this;
            DontDestroyOnLoad(gameObject);
        }
        else
        {
            Destroy(gameObject);
        }
    }

    void Start()
    {
        // Clear all persistent memory and world log when starting a new run
        WorldLog.log.Clear();

        foreach (var npc in Resources.FindObjectsOfTypeAll<NPCData>())
        {
            npc.memory.Clear();
        }

        Debug.Log("[LLMClient] Cleared all NPC memory and world log at startup.");
    }


    // Convenience wrapper: callers can pass a list of Message objects.
    // This starts a coroutine under the hood.
    public void SendMessages(List<Message> messages, Action<string> onComplete = null)
    {
        if (messages == null || messages.Count == 0)
        {
            onComplete?.Invoke("(no messages)");
            return;
        }

        StartCoroutine(SendMessagesCoroutine(messages, onComplete));
    }

    // Convenience overload for a single user/system message
    public void SendSingleMessage(string role, string content, Action<string> onComplete = null)
    {
        var list = new List<Message> { new Message(role, content) };
        SendMessages(list, onComplete);
    }

    // The coroutine that builds JSON, sends the request and parses the assistant content.
    private bool isSending = false;
    private IEnumerator SendMessagesCoroutine(List<Message> messages, Action<string> onComplete = null)
    {
        if (isSending) yield break;
        isSending = true;
        // safety checks
        if (APIConfig.Instance == null)
        {
            string err = "APIConfig.Instance is null — configure API keys and endpoints first.";
            Debug.LogError(err);
            OnRequestFailed?.Invoke(err);
            onComplete?.Invoke(null);
            yield break;
        }

        if (string.IsNullOrEmpty(endpoint) || string.IsNullOrEmpty(apiKey))
        {
            string err = "LLMClient: endpoint or apiKey is empty. Please set APIConfig at startup.";
            Debug.LogError(err);
            OnRequestFailed?.Invoke(err);
            onComplete?.Invoke(null);
            yield break;
        }

        // Build request object (use serializable structures to avoid manual string escaping)
        ChatMessage[] chatMessages = new ChatMessage[messages.Count];
        for (int i = 0; i < messages.Count; i++)
        {
            chatMessages[i] = new ChatMessage(messages[i].role ?? "user", messages[i].content ?? "");
        }

        ChatRequest requestBody = new ChatRequest
        {
            model = model,
            messages = chatMessages
        };

        string json = JsonUtility.ToJson(requestBody);
        Debug.Log($"LLMClient -> Sending request to {endpoint}. Payload:\n{json}");

        byte[] bodyRaw = System.Text.Encoding.UTF8.GetBytes(json);

        using (UnityWebRequest request = new UnityWebRequest(endpoint, "POST"))
        {
            request.uploadHandler = new UploadHandlerRaw(bodyRaw);
            request.downloadHandler = new DownloadHandlerBuffer();
            request.SetRequestHeader("Content-Type", "application/json");
            request.SetRequestHeader("Authorization", "Bearer " + apiKey);

            yield return request.SendWebRequest();

            // --- handle rate limit early ---
            if (request.responseCode == 429)
            {
                string retryAfter = request.GetResponseHeader("Retry-After");
                int wait = !string.IsNullOrEmpty(retryAfter) ? int.Parse(retryAfter) : 10;

                Debug.LogWarning($"[LLMClient] Rate limit hit. Waiting {wait}s before allowing new requests.");
                OnRateLimit?.Invoke(wait, "Rate limit exceeded");

                // Important: unlock sending so future attempts can work
                isSending = false;
                yield break;
            }

            if (request.result == UnityWebRequest.Result.Success)
            {
                string responseText = request.downloadHandler.text;
                Debug.Log("LLMClient -> Raw API response: " + Truncate(responseText, 2000));

                try
                {
                    RootResponse root = JsonUtility.FromJson<RootResponse>(responseText);
                    if (root != null && root.choices != null && root.choices.Length > 0 && root.choices[0].message != null)
                    {
                        string assistantContent = root.choices[0].message.content ?? "";
                        assistantContent = UnescapeJsonString(assistantContent);
                        onComplete?.Invoke(assistantContent);
                        isSending = false;
                        yield break;
                    }
                }
                catch (Exception ex)
                {
                    Debug.LogWarning("LLMClient: JsonUtility.FromJson failed, will fallback. " + ex.Message);
                }

                string fallback = ExtractContentWithRegex(responseText);
                if (!string.IsNullOrEmpty(fallback))
                {
                    onComplete?.Invoke(fallback);
                    isSending = false;
                    yield break;
                }

                string errMsg = "LLMClient: could not parse assistant content from response.";
                Debug.LogWarning(errMsg);
                onComplete?.Invoke(responseText);
            }
            else
            {
                string serverText = request.downloadHandler != null ? request.downloadHandler.text : "";
                string err = $"LLM request failed: {request.error}\nResponse body: {Truncate(serverText, 2000)}";
                Debug.LogError(err);
                OnRequestFailed?.Invoke(err);
                onComplete?.Invoke(null);
            }

            // Always unlock at the end
            isSending = false;

        }
        isSending = false;
    }

    // --- Helpers ---

    // Safely unescape common JSON escapes inside the extracted content
    private string UnescapeJsonString(string s)
    {
        if (string.IsNullOrEmpty(s)) return s;
        return s.Replace("\\n", "\n").Replace("\\r", "\r").Replace("\\t", "\t").Replace("\\\"", "\"").Replace("\\\\", "\\");
    }

    // Regex fallback: grabs the first "content":"..." value (handles inner escaped quotes).
    private string ExtractContentWithRegex(string json)
    {
        // This regex will capture the first content string. It's tolerant but not perfect.
        try
        {
            var m = Regex.Match(json, "\"content\"\\s*:\\s*\"([\\s\\S]*?)\"", RegexOptions.Multiline);
            if (m.Success && m.Groups.Count > 1)
            {
                string raw = m.Groups[1].Value;
                return UnescapeJsonString(raw);
            }
        }
        catch (Exception ex)
        {
            Debug.LogWarning("LLMClient: regex extraction failed: " + ex.Message);
        }
        return null;
    }

    // Utility for log truncation
    private string Truncate(string s, int max)
    {
        if (string.IsNullOrEmpty(s)) return s;
        return s.Length <= max ? s : s.Substring(0, max) + "...(truncated)";
    }
}
